###############################################################################
# Purpose: Create figures for diabetes and mortality paper
# Written by: Hunter Green
# R version: 4.4.0
###############################################################################

# Options
rm(list = ls())
cat("\014")

# Toggle for whether David is working on this
David <- F

# Set directories
if (David == T){ # David add your folder paths here
  table_dir <- paste0("/Users/dcflood/Library/CloudStorage/Dropbox-UniversityofMichigan/David Flood/HPACC/Aging projects/HRS diabetes mortality/Tables/")
  figure_dir <- paste0("/Users/dcflood/Library/CloudStorage/Dropbox-UniversityofMichigan/David Flood/HPACC/Aging projects/HRS diabetes mortality/Figures//")
} else {
  table_dir <- paste0("/Users/dcflood/Library/CloudStorage/Dropbox-UniversityofMichigan/David Flood/HPACC/Aging projects/HRS diabetes mortality/Tables/")
  figure_dir <- paste0("/Users/dcflood/Library/CloudStorage/Dropbox-UniversityofMichigan/David Flood/HPACC/Aging projects/HRS diabetes mortality/Figures/")
}

# Save date
date <- gsub("-", "_", Sys.Date())

# Load libraries
library(data.table)
library(ggplot2)
library(ggthemes)
library(readxl)

# Set default theme for graphics
theme_set(theme_classic())


###############################################################################
# Load data
###############################################################################
# Figure 1
fig1_data <- as.data.table(read_excel(paste0(table_dir, "Figure1_data.xlsx")))

# Figure 2
fig2_data <- as.data.table(read_excel(paste0(table_dir, "Figure2_data.xlsx")))


###############################################################################
# Prevalence Figure - Age and Cohort
###############################################################################
# Create variables
fig1_data[, agenum := fcase(age == "51-59", 1,
                            age == "60-69", 2,
                            age == "70+", 3)]
fig1_data[, agenum := ifelse(study == "China (CHARLS)", agenum - 0.05,
                      ifelse(study == "England (ELSA)", agenum + 0.05, agenum))]

# Factorize group/color variable
fig1_data$study <- factor(fig1_data$study,
                          levels = c("China (CHARLS)", "England (ELSA)", "Mexico (MHAS)", "South Africa (HAALSI)", "United States (HRS)"))
fig1_data$age <- as.factor(fig1_data$age)

# Numerize variable
fig1_data$pct <- as.numeric(fig1_data$pct)
fig1_data$pct_ll <- as.numeric(fig1_data$pct_ll)
fig1_data$pct_ul <- as.numeric(fig1_data$pct_ul)

# Create plot
fig_1 <- ggplot(data = fig1_data, aes(x = agenum, y = pct)) +
  # Plot points
  geom_point(aes(color = study)) +
  # Plot error bars
  geom_linerange(aes(ymin = pct_ll, ymax = pct_ul, color = study),
                 linewidth = 0.6, alpha = 0.8) +
  # Plot lines to connect points
  geom_line(aes(color = study), linewidth = 0.7, alpha = 0.5, show.legend = FALSE) +
  # Annotations at end of lines
  geom_text(data = subset(fig1_data, age == "70+" & study == "China (CHARLS)"),
            aes(label = "China (2011-12)", color = study),
            size = 3, hjust = 0, nudge_x = 0.07, nudge_y = -1.5) +
  geom_text(data = subset(fig1_data, age == "70+" & study == "England (ELSA)"),
            aes(label = "England (2012-13)", color = study),
            size = 3, hjust = 0, nudge_x = 0.05) +
  geom_text(data = subset(fig1_data, age == "70+" & study == "Mexico (MHAS)"),
            aes(label = "Mexico (2012-13)", color = study),
            size = 3, hjust = 0, nudge_x = 0.05) +
  geom_text(data = subset(fig1_data, age == "70+" & study == "South Africa (HAALSI)"),
            aes(label = "South Africa \n(2014-15)", color = study),
            size = 3, hjust = 0, nudge_x = 0.05, nudge_y = -1.5) +
  geom_text(data = subset(fig1_data, age == "70+" & study == "United States (HRS)"),
            aes(label = "United States \n(2010-13)", color = study),
            size = 3, hjust = 0, nudge_x = 0.05, nudge_y = -1.5) +
  # Axis labels
  labs(x = "Age (years)",
       y = "Prevalence (%)") +
  # Theme
  theme(# Title
        title = element_text(size=24),
        # Axis
        axis.title = element_text(size = 12),
        axis.title.x = element_text(margin = margin(t = 10)),
        axis.title.y = element_text(margin = margin(r = 10)),
        axis.text.x = element_text(size = 10, color = "black"),
        axis.text.y = element_text(size = 10, color = "black"),
        # Legend
        legend.position = "none",
        # Panel
        panel.grid.major.y = element_line(colour = "grey90"),
        panel.spacing = unit(2, "lines"),
        # Margin
        plot.margin=unit(c(t = 15, r = 15, b = 15, l = 15), "points"),
        # Strip
        strip.background = element_blank()) +
  # Axis scale
  scale_y_continuous(breaks = seq(from = 0, to = 50, by = 10), limits = c(0, 55), expand = c(0, 0)) +
  scale_x_continuous(breaks = seq(from = 1, to = 3, by = 1), limits = c(0.5, 3.75), expand = c(0, 0),
                     labels = c("51-59", "60-69", "70+")) +
  scale_color_stata()
  
# Display plot
fig_1
  
# Save plot
ggsave(paste0(figure_dir, "Figure1_prev_", date, ".png"), plot = fig_1, height = 4.5, width = 6)

########### Updated figure 2
# Create bar chart
fig_1 <- ggplot(data = fig1_data, aes(x = age, y = pct, fill = study)) +
  # Plot bars
  geom_bar(stat = "identity", position = position_dodge(width = 0.9)) +
  # Add error bars as vertical lines
  geom_errorbar(aes(ymin = pct_ll, ymax = pct_ul), 
                position = position_dodge(width = 0.9), width = 0, linewidth = 0.6, alpha = 0.8) +
  # Axis labels
  labs(x = "Age (years)",
       y = "Prevalence (%)",
       fill = NULL) +  # Remove legend title
  # Theme
  theme(# Title
    title = element_text(size=24),
    # Axis
    axis.title = element_text(size = 12),
    axis.title.x = element_text(margin = margin(t = 10)),
    axis.title.y = element_text(margin = margin(r = 10)),
    axis.text.x = element_text(size = 10, color = "black"),
    axis.text.y = element_text(size = 10, color = "black"),
    # Legend
    legend.position = "right",
    legend.text = element_text(size = 8),  # Smaller legend text
    legend.key.size = unit(0.5, "cm"),  # Reduce legend key size
    # Panel
    panel.grid.major.y = element_line(colour = "grey90"),
    panel.spacing = unit(2, "lines"),
    # Margin
    plot.margin=unit(c(t = 15, r = 15, b = 15, l = 15), "points"),
    # Strip
    strip.background = element_blank()) +
  # Axis scale
  scale_y_continuous(breaks = seq(from = 0, to = 50, by = 10), limits = c(0, 55), expand = c(0, 0)) +
  scale_fill_stata()  # Adjust color palette for fill

# Display plot
fig_1

# Save plot
ggsave(paste0(figure_dir, "Figure1_prev_bar_corrected_", date, ".png"), plot = fig_1, height = 4.5, width = 6)
###############################################################################
# Mortality Rate Figure 
###############################################################################
# Factorize variables
fig2_data$study <- factor(fig2_data$study,
                          levels = c("China (CHARLS)", "England (ELSA)", "Mexico (MHAS)", "South Africa (HAALSI)", "United States (HRS)"))
fig2_data$diabetes <- factor(fig2_data$diabetes,
                             levels = c("Without diabetes", "Diabetes"))

# Numerize variable
fig2_data$rate <- as.numeric(fig2_data$rate)
fig2_data$rate_ll <- as.numeric(fig2_data$rate_ll)
fig2_data$rate_ul <- as.numeric(fig2_data$rate_ul)

# Create plot
fig_2 <- ggplot(data = fig2_data, aes(x = study, y = rate, fill = diabetes)) +
  # Plot points
  geom_bar(position = "dodge", stat = "identity", color = "black", linewidth = 0.3, width = 0.8) +
  # Plot error bars
  geom_linerange(aes(ymin = rate_ll, ymax = rate_ul),
                 linewidth = 0.6, position = position_dodge(width = 0.8), color = "black") +
  # Set data labels
  geom_text(data = subset(fig2_data, diabetes == "Without diabetes"),
            aes(label = format(round(rate, digits = 1), nsmall = 1)),
            size = 3, nudge_x = -0.375, nudge_y = 2) +
  geom_text(data = subset(fig2_data, diabetes == "Diabetes"),
            aes(label = format(round(rate, digits = 1), nsmall = 1)),
            size = 3, nudge_x = 0.375, nudge_y = 2) +
  # Axis labels
  labs(x = "",
       y = "Mortality rate (per 1,000 person-years)") +
  theme(# Title
    title = element_text(size = 24),
    # Axis
    axis.title = element_text(size = 12),
    axis.title.x = element_text(margin = margin(t = 10)),
    axis.title.y = element_text(margin = margin(r = 10)),
    axis.text.x = element_text(size = 10, color = "black"),
    axis.text.y = element_text(size = 10, color = "black"),
    # Legend
    legend.position = c(0.15, 0.90),
    legend.background = element_rect(colour = 'black', fill = 'white', linetype = 'solid', linewidth = 0.3),
    legend.margin = margin(c(t = 4, r = 2, b = 4, l = 4)),
    legend.box.margin = margin(c(t = 0, r = 0, b = 0, l = 0)),
    legend.text = element_text(margin = margin(t = 2, r = 2, b = 2, l = 2, unit = "pt")),
    legend.title = element_blank(),
    legend.spacing.x = unit(0.0, 'pt'),
    legend.spacing.y = unit(0.0, 'pt'),
    legend.key.size = unit(10, "pt"),
    # Panel
    panel.grid.major.y = element_line(colour = "grey90"),
    panel.spacing = unit(2, "lines"),
    # Margin
    plot.margin=unit(c(t = 15, r = 15, b = 15, l = 15), "points"),
    # Strip
    strip.background = element_blank()) +
  # Axis scale
  scale_y_continuous(breaks = seq(from = 0, to = 70, by = 10), limits = c(0, 75), expand = c(0, 0)) +
  scale_x_discrete(breaks = c("China (CHARLS)", "England (ELSA)", "Mexico (MHAS)", "South Africa (HAALSI)", "United States (HRS)"),
                   labels = c("China\n (CHARLS)", "England\n (ELSA)", "Mexico\n (MHAS)", "South Africa\n (HAALSI)", "United States\n (HRS)")) +
  scale_fill_stata()

# Display plot
fig_2

# Save plot
ggsave(paste0(figure_dir, "Figure2_mortalityrate_", date, ".png"), plot = fig_2, height = 4.5, width = 6)

